'''
@Version: 0.0.2
@Author: ider
@Date: 2019-12-27 11:29:07
@LastEditors  : ider
@LastEditTime : 2020-02-06 14:24:33
@Description: 计算学科内的小世界图
'''


import os
import csv
import json
import pymongo
import codecs
import graph_tool.all as gt
import numpy as np

import logging
from config import MONGO_URL
from DataGet import WikiLinkYear


class Graph:
    def __init__(self,edges):
        # 无向图
        self.graph = gt.Graph(directed=False)
        # 添加所有的边
        if isinstance(edges,list):
            self.graph.add_edge_list(edges,hashed=True)
        else:
            self.graph.add_edge_list(edges(),hashed=True)
        # 去除重复边
        gt.remove_parallel_edges(self.graph)
        # 获得联通子图
        self.component_graph = gt.extract_largest_component(self.graph)
        
    def get_avg_path(self):
        all_sp = gt.shortest_distance(self.component_graph)
#         vertex_avgs = gt.vertex_average(self.component_graph, all_sp)
#         return float(np.mean(vertex_avgs[0]))
        return float(np.sum([np.sum(i.get_array()) for i in all_sp])/(self.component_graph.num_vertices()**2-self.component_graph.num_vertices()))
    
    
    def get_clustering_coefficient(self):
        clust = gt.local_clustering(self.component_graph)
        vertex_avgs =gt.vertex_average(self.component_graph, clust)
        return vertex_avgs[0]
    
    def get_num_edges(self):
        return self.component_graph.num_edges()
    
    def get_num_vertices(self):
        return self.component_graph.num_vertices()


def main(desyear=None,level=2):
    Table = pymongo.MongoClient(MONGO_URL).small_world[f'mini_world_level{level}']
    if not desyear:
        years = [i for i in range(2007,2021)]
        del(years[years.index(2012)])
        del(years[years.index(2010)])
    else:
        years = [desyear,]
    for year in years:
        logging.info(year)
        wly = WikiLinkYear(year)
        
        # 每年对应的 大类 ids
        _categorys_ids_dict,categorys_ids_set = wly.get_article_ids(level=level)

        # 从文件中加载所有的 edge
        category_edgs_dict = {}
        for key,value in wly.iterate_edges():
            for cat,cat_id_set in categorys_ids_set.items():
                if key in cat_id_set and value in cat_id_set:
                    category_edgs_dict.setdefault(cat,[])
                    category_edgs_dict[cat].append([key,value])
        

        for cate_name ,edges in category_edgs_dict.items():
            logging.info(f'edges count: {cate_name},{len(edges)}')
            graph = Graph(edges)

            avg_path = graph.get_avg_path()
            clustering_coefficient = graph.get_clustering_coefficient()
            num_edges = graph.get_num_edges()
            num_vertices = graph.get_num_vertices()
            print(year,cate_name,avg_path,clustering_coefficient)

            Table.update_one({'category_name':cate_name,'year':year},{'$set':{'avg_path':avg_path,
                                                                            'clustering_coefficient':clustering_coefficient,
                                                                            'num_edges':num_edges,
                                                                            'num_vertices':num_vertices}},upsert=True)

if __name__ == "__main__":
    main()